from pyspark.sql import SparkSession
import pandas as pd
from gai.v2.spark.transformer import DayRectifier
from gai.v2.spark.transformer import FeatureRetriever

spark = SparkSession \
    .builder \
    .appName('统计') \
    .enableHiveSupport() \
    .getOrCreate()


gid_list = [
    'ANDROID-e5315edf4bf94ed8b6ca6fda3ad2f7eb',
]

GIDpd = pd.DataFrame(gid_list, columns=['gid'])
df_tmp = spark.createDataFrame(GIDpd)
df_tmp.show()

df_tmp_pd = df_tmp.toPandas()
df_tmp_pd['day'] = '20200601'
df_tmp_pd['day'] = df_tmp_pd['day'].apply(lambda x: str(x))
df_tmp_pd['matched_gid'] = df_tmp_pd['gid']
df_tmp_df = spark.createDataFrame(df_tmp_pd)
df_tmp_df.show()


day_rectifier = DayRectifier(inputCol='day', outputCol='rectified_day')
rectified_day_df = day_rectifier.transform(df_tmp_df)
print(rectified_day_df.show())

cols = [
    'ft_usertags'
    , 'ft_installapp_ls'
    , 'ft_brand'
    , 'ft_phone_model'
]


feature_retriever = FeatureRetriever(inputIdCol="matched_gid",
                                     inputDayCol="rectified_day",
                                     outputFeatureCols=cols,
                                     extraParams={'span.in.months': 1})
retrieved_feature_df = feature_retriever.transform(rectified_day_df)
retrieved_feature_df.show()
